package com.uziot.bucket.guava.limit;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.util.concurrent.RateLimiter;
import com.uziot.bucket.guava.pojo.HttpContextUtil;
import com.uziot.bucket.guava.pojo.ResultCode;
import com.uziot.bucket.guava.pojo.ResultUtil;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 功能描述: <br>
 * 接口限流切面类
 * <br>
 * 使用原理：
 * RateLimiter.create(limitNum) 可指定每秒放入的数量
 * 如果超过指定的数量，是无法放入的
 * 使用这个特性，可以进行限流操作
 *
 * @author shidt
 * @date 2020-03-09 0:00
 */
@Slf4j
@Aspect
@Order(5)
@Component
@SuppressWarnings("UnstableApiUsage")
public class LimitGuavaAspect {
    /**
     * 用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
     */
    private final ConcurrentHashMap<String, RateLimiter> map = new ConcurrentHashMap<>();
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    @Pointcut("@annotation(com.uziot.bucket.guava.limit.LimitGuava)")
    public void serviceLimit() {
    }

    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) throws NoSuchMethodException {
        Object obj = null;
        //获取拦截的方法名
        Signature sig = joinPoint.getSignature();
        //获取拦截的方法名
        MethodSignature methodSignature = (MethodSignature) sig;
        //返回被织入增加处理目标对象
        Object target = joinPoint.getTarget();
        //为了获取注解信息
        Method currentMethod = target.getClass().getMethod(methodSignature.getName(), methodSignature.getParameterTypes());
        //获取注解信息
        LimitGuava annotation = currentMethod.getAnnotation(LimitGuava.class);
        //获取注解每秒加入桶中的token
        double limitNum = annotation.limitSecondNum();
        // 注解所在方法名区分不同的限流策略
        String functionName = methodSignature.getName();

        //获取rateLimiter
        if (!map.containsKey(functionName)) {
            log.info("GuavaLimit，方法名称:{},限制次数为:{}", functionName, limitNum);
            map.put(functionName, RateLimiter.create(limitNum));
        }
        RateLimiter rateLimiter = map.get(functionName);
        try {
            if (rateLimiter.tryAcquire()) {
                //执行方法
                obj = joinPoint.proceed();
            } else {
                //拒绝了请求（服务降级）
                String result = OBJECT_MAPPER.writeValueAsString(ResultUtil.error(ResultCode.ERROR, "接口访问超出频率限制！"));
                log.error("接口访问超出频率限制,每秒限制频率每秒:【{}】次,result:{}", limitNum, result);
                //调用返回方法，适合非包装类返回，最好是String返回控制器
                outErrorResult(result);
            }
        } catch (Throwable throwable) {
            throwable.printStackTrace();
        }
        return obj;
    }

    /**
     * 将结果返回
     *
     * @param result res
     */
    public void outErrorResult(String result) {
        HttpServletResponse response = HttpContextUtil.getHttpServletResponse();
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        try (ServletOutputStream outputStream = response.getOutputStream()) {
            outputStream.write(result.getBytes(StandardCharsets.UTF_8));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    static {
        OBJECT_MAPPER.setSerializationInclusion(JsonInclude.Include.NON_NULL);
    }
}